## Bake Animation
## by Stefan Ihringer, stefan@bildfehler.de
##
## Python tool script example based on the "Bake Animation" script that ships with Fusion.
## lessons:
## - retrieving inputs from a tool
## - retrieving attributes of an input
## - accessing inputs at arbitrary frames
## - creating bezier spline animation
## - locking and undo of composition
##
## version 1.1 (2012-06-23): fix range (last frame wasn't baked!), added progress output to console in case
## baking takes a long time (for example when baking a probe)


def get_valid_inputs(t):
	"This function retrieves all (valid) animated inputs from a given tool as a dict (ID => input object)."
	
	# We can't bake DT_Image or DT_Mask (and possibly others...)
	unbakeable = ["Image", "Mask", "Particles", "DataType3D"]

	inputlist = t.GetInputList()
	result = {}
	for inp in inputlist.itervalues():
		if inp.GetConnectedOutput() is not None or inp.GetExpression() is not None:
			if inp.GetAttrs()["INPS_DataType"] not in unbakeable:
				result[inp.ID] = inp
	#print(result)
	return result


def bake_inputs(inputlist, timestart, timeend):
	"This function bakes the given inputs between timestart and timeend."
	
	global composition
	global tool
	bakedkeys = {}
	if len(inputlist) == 0:
		return

	composition.Lock()
	try:
		composition.StartUndo("Bake Animation")
		composition.GetFrameList()[1].SwitchMainView('ConsoleView') 
		print "recording keys..."

		# step through frame range and record keys
		for f in range(min(timestart, timeend), max(timestart, timeend) + 1):
			if f == timestart or f == timeend or f % 10 == 0:
				print "%d" % f
			for inp in inputlist:
				try:
					bakedkeys[inp.ID][f] = inp[f]
				except KeyError:
					bakedkeys[inp.ID] = {}
					bakedkeys[inp.ID][f] = inp[f]

		print "baking keys of:"
		for inp in inputlist:
			print "- " + inp.Name
			if tool[inp.ID].GetExpression() is not None:
				tool[inp.ID].SetExpression(None)
			# overwrite inputs with fresh modifiers. These splines will
			# already contain a key frame at the current time that has
			# to be removed if it isn't overwritten during the next loop.
			if inp.GetAttrs()["INPS_DataType"] == "Point":
				tool.AddModifier(inp.ID, "PolyPath")
			else:
				tool.AddModifier(inp.ID, "BezierSpline")

		# write recorded keys back
		currenttime = comp.CurrentTime
		for f in range(min(timestart, timeend), max(timestart, timeend) + 1):
			if f == timestart or f == timeend or f % 10 == 0:
				print "%d" % f
			for inp in inputlist:
				tool[inp.ID][f] = bakedkeys[inp.ID][f]
				# delete the auto-generated keyframe
				if currenttime is not None and f != currenttime:
					tool[inp.ID][currenttime] = None
			currenttime = None
			
		composition.EndUndo(True)
		print "done"

	finally:
		composition.Unlock()
	
	
	
## main ##
try:
	
	try:
		if tool is None:
			tool = comp.ActiveTool
	except NameError:
		tool = comp.ActiveTool

	if tool is None:
		raise Exception("This is a tool script, you must select a tool in the flow to run this script.")
		
	# get all animated inputs
	inputs = get_valid_inputs(tool)

	if len(inputs) == 0:
		raise Exception("This tool has no inputs which can be baked by this script.")
		
	# build user dialog
	dialog = {}
	i = 1
	for inp in inputs.itervalues():
		# add a small explanation on how each input is animated
		label = inp.Name
		if inp.GetExpression() is not None:
			label = label + " (expression)"
		else:
			label = label + " (" + inp.GetConnectedOutput().Name + " from " + inp.GetConnectedOutput().GetTool().ID + ")"
		dialog[i] = {1: inp.ID, 2: "Checkbox", "Name": label}
		i = i + 1

	compattrs = composition.GetAttrs()	
	dialog[i]   = {1: "from", 2: "Slider", "Name": "From:", "Default": compattrs["COMPN_GlobalStart"], "Min": compattrs["COMPN_GlobalStart"], "Max": compattrs["COMPN_GlobalEnd"], "Integer": True }
	dialog[i+1] = {1: "to",   2: "Slider", "Name": "To:",   "Default": compattrs["COMPN_GlobalEnd"],   "Min": compattrs["COMPN_GlobalStart"], "Max": compattrs["COMPN_GlobalEnd"], "Integer": True }

	# show dialog box
	ret = composition.AskUser("Select inputs to be baked:", dialog)

	if ret is not None:
		#print(ret)
		# remove unselected inputs from list.
		# ret will look something like this: {"from": 1, "to": 100, "Input1": 1.0, "Input2": 0.0}
		for k, v in ret.iteritems():
			try:
				if v == 0:
					del inputs[k]
			except (ValueError, KeyError):
				# in case this loop tries to remove items that are not from the checkboxes...
				pass
				
		bake_inputs(inputs.values(), int(ret["from"]), int(ret["to"]))

except Exception as e:
	print e
	
# fin